import gc
import json

import cv2
import numpy as np
import torch

from utils.rapid import rapid_recognize

# 检查是否有可用的 GPU，如果有则使用 GPU 进行推理，否则使用 CPU
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'


def detect_objects(img, model, ocr_engine, table_engine):
    """
    使用 YOLO 和 RapidTable 对输入图片进行检测，识别表格并提取信息。

    :param img: 输入的 PIL 图像对象
    :param model: YOLO 模型
    :param ocr_engine: 光学字符识别引擎
    :param table_engine: 表格识别引擎
    :return: 包含检测结果的 JSON 字符串
    """
    # 使用 YOLO 模型对输入图像进行预测，设置置信度阈值为 0.6
    with torch.no_grad():
        results = model.predict(img, conf=0.6, device=device)

    result_list = []

    try:
        # 检查 results 是否有效
        if results is None or len(results) == 0:
            return json.dumps([])

        # 遍历每个检测结果
        for result_idx, result in enumerate(results):
            # 检查 boxes 是否有效
            if not hasattr(result, 'boxes') or result.boxes is None or len(result.boxes) == 0:
                continue

            boxes = result.boxes

            # 遍历每个边界框
            for box_idx, box in enumerate(boxes):
                try:
                    # 获取边界框信息
                    confidence = box.conf[0].item()
                    # 获取边界框坐标（ndarray类型）
                    coords_ndarray = box.xyxy[0].cpu().numpy()
                    # 转换为整数列表
                    coords_list = coords_ndarray.astype(int).tolist()
                    x1, y1, x2, y2 = coords_list

                    # 表格区域（略微扩大边界）
                    table_region = {
                        'x1': max(0, x1 - 20),
                        'y1': max(0, y1 - 20),
                        'x2': min(img.width, x2 + 20),
                        'y2': min(img.height, y2 + 20)
                    }

                    # 表格名称区域（表格上方区域）
                    table_height = y2 - y1
                    name_region = {
                        'x1': max(0, x1 - 20),
                        'y1': max(0, y1 - int(table_height * 0.3)),
                        'x2': min(img.width, x2 + 20),
                        'y2': y1
                    }

                    # 识别表格名称
                    table_name = recognize_table_name(img, name_region, ocr_engine)

                    # 提取表格内容
                    cropped_img = img.crop((table_region['x1'], table_region['y1'],
                                            table_region['x2'], table_region['y2']))
                    cropped_img_cv2 = pil_to_cv2(cropped_img)
                    html_result = rapid_recognize(cropped_img_cv2, ocr_engine, table_engine)

                    # 添加结果
                    result_list.append({
                        'confidence': confidence,
                        'bbox': coords_list,
                        'html_result': html_result,
                        'table_name': table_name
                    })

                except Exception as e:
                    print(f"处理第{result_idx}个结果的第{box_idx}个边界框时出错: {e}")
                    continue

            # 释放当前结果的资源
            del result, boxes
            torch.cuda.empty_cache()

        # 释放所有结果的资源
        del results
        torch.cuda.empty_cache()
        print(result_list)
        return json.dumps(result_list)

    except Exception as e:
        print(f"检测过程中出错: {e}")
        return json.dumps([{"error": f"Error during detection: {str(e)}"}])

    finally:
        gc.collect()


def pil_to_cv2(pil_img):
    """将PIL图像转换为OpenCV格式(BGR)"""
    img_np = np.array(pil_img)
    return cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)


def recognize_table_name(img, region, ocr_engine):
    """
    识别表格名称区域的文本

    :param img: 原始PIL图像
    :param region: 表格名称区域坐标 {'x1', 'y1', 'x2', 'y2'}
    :param ocr_engine: OCR引擎
    :return: 识别出的表格名称
    """
    try:
        # 裁剪名称区域
        name_region_img = img.crop((region['x1'], region['y1'], region['x2'], region['y2']))
        name_region_cv2 = pil_to_cv2(name_region_img)

        # 使用OCR识别文本
        ocr_output = ocr_engine(name_region_cv2, return_word_box=True)

        # 检查OCR输出是否有效
        if (ocr_output is not None and hasattr(ocr_output, 'txts') and ocr_output.txts
                and len(ocr_output.txts) > 0):
            # 提取文本并连接
            return ' '.join(ocr_output.txts).strip() or "未识别到表格名称"

        return "未识别到表格名称"

    except Exception as e:
        print(f"识别表格名称时出错: {e}")
        return "识别表格名称失败"
